package com.example.demo.loader;

import com.sun.tools.javac.file.JavacFileManager;
import org.springframework.boot.loader.jar.JarFile;

import javax.tools.*;
import javax.tools.JavaFileObject.Kind;
import java.io.*;
import java.net.URI;
import java.net.URL;
import java.nio.CharBuffer;
import java.util.*;
import java.util.jar.JarEntry;
import java.util.stream.Collectors;

/**
 * 内存Java文件管理器
 * 用于加载引用的lib下面的依赖资源，jar包内的class
 */
public class MemoryJavaFileManager extends ForwardingJavaFileManager<JavaFileManager> {

	final Map<String, byte[]> classBytes = new HashMap<>();

	final Map<String, List<JavaFileObject>> classObjectPackageMap = new HashMap<>();

	private JavacFileManager javaFileManager;

	/**
	 * key 包名 value javaObj 主要给jdk编译class的时候找依赖class用
	 */
	public final static Map<String, List<JavaFileObject>> CLASS_OBJECT_PACKAGE_MAP = new HashMap<>();

	private final static Object LOCK = new Object();

	private boolean isInit = false;

	/**
	 * 初始化获取jar项目引入的第三方jar包class
	 */
	public void init() {
		try {
			String jarBaseFile = MemoryClassLoader.getPath();
			JarFile jarFile = new JarFile(new File(jarBaseFile));
			// 解决找不到service、bean、util等自己写的class
			insCacheClass(jarFile, true);
			// 获取所引入的jar包
			List<JarEntry> entries = jarFile.stream().filter(jarEntry -> jarEntry.getName().endsWith(".jar")).collect(Collectors.toList());
			JarFile libTempJarFile;
			// 遍历每一个引用的jar包
			for (JarEntry entry : entries) {
				libTempJarFile = jarFile.getNestedJarFile(jarFile.getEntry(entry.getName()));
				if (libTempJarFile.getName().contains("tools.jar")) {
					continue;
				}
				insCacheClass(libTempJarFile, false);
			}
		} catch (Exception e) {
			e.printStackTrace();
		}
		isInit = true;
	}

	MemoryJavaFileManager(JavaFileManager fileManager) {
		super(fileManager);
		this.javaFileManager = (JavacFileManager) fileManager;
	}

	public Map<String, byte[]> getClassBytes() {
		return new HashMap<>(this.classBytes);
	}

	@Override
	public void flush() throws IOException {
	}

	@Override
	public void close() throws IOException {
		classBytes.clear();
	}

	/**
	 * 从jar内遍历class并缓存
	 * @param jarFile jar文件
	 * @param isProjectJar 是否本项目jar
	 */
	private void insCacheClass(JarFile jarFile, Boolean isProjectJar) {
		Enumeration<JarEntry> tempEntriesEnum = jarFile.entries();
		String packageName;
		List<JavaFileObject> onePackageJavaFiles;
		// 相关的class文件进行缓存
		try {
			while (tempEntriesEnum.hasMoreElements()) {
				JarEntry jarEntry = tempEntriesEnum.nextElement();
				String classPath = jarEntry.getName().replace("/", ".");
				if (classPath.endsWith(".class") && jarEntry.getName().lastIndexOf("/") != -1) {
					packageName = classPath.substring(0, jarEntry.getName().lastIndexOf("/"));
					// jar项目
					if (isProjectJar) {
						if (packageName.contains("BOOT-INF.classes.")) {
							packageName = packageName.replace("BOOT-INF.classes.", "");
						}
					}
					onePackageJavaFiles = CLASS_OBJECT_PACKAGE_MAP.containsKey(packageName) ? CLASS_OBJECT_PACKAGE_MAP.get(packageName) : new ArrayList<>();
					onePackageJavaFiles.add(new MemoryJavaClassObject(jarEntry.getName().replace("/", ".").replace(".class", ""),
							new URL(jarFile.getUrl(), jarEntry.getName()), javaFileManager));
					CLASS_OBJECT_PACKAGE_MAP.put(packageName, onePackageJavaFiles);
				}
			}
		} catch (Exception e) {
			e.printStackTrace();
		}
	}

	/**
	 * 获取依赖包
	 * @param packageName 限定名
	 * @return /
	 */
	public List<JavaFileObject> getLibJarsOptions(String packageName) {
		synchronized (LOCK) {
			if (!isInit) {
				init();
			}
		}
		return CLASS_OBJECT_PACKAGE_MAP.get(packageName);
	}

	@Override
	public Iterable<JavaFileObject> list(Location location,
	                                     String packageName,
	                                     Set<Kind> kinds,
	                                     boolean recurse)
			throws IOException {
		// 判断是否是jar包形式运行项目
		if ("CLASS_PATH".equals(location.getName()) && MemoryClassLoader.isJar()) {
			// packageName => Java文件import的类限定名
			List<JavaFileObject> result = getLibJarsOptions(packageName);
			if (result != null) {
				return result;
			}
		}
		Iterable<JavaFileObject> it = super.list(location, packageName, kinds, recurse);
		if (kinds.contains(Kind.CLASS)) {
			final List<JavaFileObject> javaFileObjectList = classObjectPackageMap.get(packageName);
			if (javaFileObjectList != null) {
				if (it != null) {
					for (JavaFileObject javaFileObject : it) {
						javaFileObjectList.add(javaFileObject);
					}
				}
				return javaFileObjectList;
			} else {
				return it;
			}
		} else {
			return it;
		}
	}

	@Override
	public String inferBinaryName(Location location, JavaFileObject file) {
		if (file instanceof MemoryInputJavaClassObject) {
			return ((MemoryInputJavaClassObject) file).inferBinaryName();
		}
		return super.inferBinaryName(location, file);
	}

	@Override
	public JavaFileObject getJavaFileForOutput(Location location, String className, Kind kind,
	                                           FileObject sibling) throws IOException {
		if (kind == Kind.CLASS) {
			return new MemoryOutputJavaClassObject(className);
		} else {
			return super.getJavaFileForOutput(location, className, kind, sibling);
		}
	}

	JavaFileObject makeStringSource(String className, final String code) {
		String classPath = className.replace('.', '/') + Kind.SOURCE.extension;
		return new SimpleJavaFileObject(URI.create("string:///" + classPath), Kind.SOURCE) {
			@Override
			public CharBuffer getCharContent(boolean ignoreEncodingErrors) {
				return CharBuffer.wrap(code);
			}
		};
	}

	void makeBinaryClass(String className, final byte[] bs) {
		JavaFileObject javaFileObject = new MemoryInputJavaClassObject(className, bs);
		String packageName = "";
		int pos = className.lastIndexOf('.');
		if (pos > 0) {
			packageName = className.substring(0, pos);
		}
		List<JavaFileObject> javaFileObjectList = classObjectPackageMap.get(packageName);
		if (javaFileObjectList == null) {
			javaFileObjectList = new LinkedList<>();
			javaFileObjectList.add(javaFileObject);
			classObjectPackageMap.put(packageName, javaFileObjectList);
		} else {
			javaFileObjectList.add(javaFileObject);
		}
	}

	class MemoryInputJavaClassObject extends SimpleJavaFileObject {
		final String className;
		final byte[] bs;

		MemoryInputJavaClassObject(String className, byte[] bs) {
			super(URI.create("string:///" + className.replace('.', '/') + Kind.CLASS.extension), Kind.CLASS);
			this.className = className;
			this.bs = bs;
		}

		@Override
		public InputStream openInputStream() {
			return new ByteArrayInputStream(bs);
		}

		public String inferBinaryName() {
			return className;
		}
	}


	class MemoryOutputJavaClassObject extends SimpleJavaFileObject {
		final String className;
		MemoryOutputJavaClassObject(String className) {
			super(URI.create("string:///" + className.replace('.', '/') + Kind.CLASS.extension), Kind.CLASS);
			this.className = className;
		}

		@Override
		public OutputStream openOutputStream() {
			return new FilterOutputStream(new ByteArrayOutputStream()) {
				@Override
				public void close() throws IOException {
					out.close();
					ByteArrayOutputStream bos = (ByteArrayOutputStream) out;
					byte[] bs = bos.toByteArray();
					classBytes.put(className, bs);
					makeBinaryClass(className, bs);
				}
			};
		}
	}
}
